import numpy as np
from PIL import Image
import joblib
import gradio as gr


# 假设模型和缩放器已经加载
model = joblib.load('models/mlp_mnist_model.pkl')
scaler = joblib.load('models/mlp_scaler.pkl')


# 图像预处理函数
def preprocess_image(input_data):
    if isinstance(input_data, dict):
        image_data = input_data.get('composite', None)
        if image_data is None:
            image_data = input_data.get('image', None)
    else:
        image_data = input_data

    if image_data is None or np.all(image_data == 0):
        print("Empty or invalid image data received.")
        return None

    # 将图像转换为灰度图像并调整大小
    pil_image = Image.fromarray(image_data).convert('L')
    pil_image = pil_image.resize((28, 28), Image.LANCZOS)

    # 颜色反转（黑色背景、白色数字）

    # 将图像数据转换为numpy数组并扁平化
    image_array = np.array(pil_image).astype(float)  # 确保转换为浮点数
    image_flattened = image_array.flatten()

    # 返回图像数组和处理后的 28x28 图像
    return image_flattened.reshape(1, -1), image_array  # 这里 reshape 成 2D 数组


# 定义预测函数
def predict_digit(image):
    if image is None:
        return {}, "", None

    # 预处理图像，并获取 28x28 图像
    processed_image, processed_image_28x28 = preprocess_image(image)

    if processed_image is None:
        return {}, "Invalid image", None

    # 标准化图像数据
    processed_image = scaler.transform(processed_image)

    # 进行预测
    prediction = model.predict(processed_image)[0]
    probabilities = model.predict_proba(processed_image)[0]

    # 创建结果字典
    results = {str(i): float(prob) for i, prob in enumerate(probabilities)}

    # 将预处理后的图像转换为可以显示的格式
    pil_image = Image.fromarray(processed_image_28x28).convert('L')

    return results, str(prediction), pil_image


# 创建 Gradio 接口
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Image(label="上传图片"),  # 将手写板替换为上传图片
    outputs=[
        gr.Label(num_top_classes=3, label="预测概率"),
        gr.Textbox(label="预测结果"),
        gr.Image(label="预处理后的28x28图像")  # 显示预处理后的图像
    ],
    title="手写数字识别",
    description="上传一张图片，然后点击预测按钮。"
)

iface.launch()
